import json
import traceback
from datetime import datetime, timezone
import pytz

from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.background import BackgroundTasks
# from models.base import LogEntry


class LogMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        # # 获取当前的北京时间
        # if ((request.url.path == '/api/v1/backend/product/list' or request.url.path == '/api/v1/chat/completions/biz'
        #         or request.url.path == '/api/v1/external/chat/completions/biz')
        #         or request.url.path == '/api/v1/smart_backend/chat/report'
        #         or request.url.path == '/api/v1/smart_backend/message/report'):
        #     return await call_next(request)
        beijing_tz = pytz.timezone('Asia/Shanghai')
        start_time_without_ms = datetime.now(beijing_tz)
        start_time = start_time_without_ms.strftime('%Y-%m-%d %H:%M:%S')
        # print(request.url.path, request.method, start_time)

        # 读取并复制原始请求体
        body = await request.body()
        request.state.body = body

        # 解析请求体参数，用于日志
        try:
            params = json.loads(body.decode()) if body else {}
        except json.JSONDecodeError:
            params = {}

        query_params = dict(request.query_params)
        client_ip = request.client.host

        # 重构请求，以便可以重新读取请求体
        async def fake_receive():
            return {'type': 'http.request', 'body': body, 'more_body': False}

        new_request = Request(request.scope, receive=fake_receive)

        try:
            response = await call_next(new_request)

            # 处理响应日志
            response_content = b''
            async for chunk in response.body_iterator:
                response_content += chunk

            if response.headers.get('Content-Type', '').startswith('application/json'):
                try:
                    response_body_str = json.dumps(json.loads(response_content.decode()),
                                                   ensure_ascii=False) if response_content else ''
                except json.JSONDecodeError:
                    response_body_str = response_content.decode() if response_content else ''
            else:
                response_body_str = response_content.decode() if response_content else ''

            end_time = datetime.now(beijing_tz)
            duration = (end_time - start_time_without_ms).total_seconds()

            log_entry = LogEntry(
                timestamp=start_time,
                level="INFO",
                module=request.url.path,
                function=request.method,
                message="API call",
                params={**params, **query_params},
                response_status=response.status_code,
                response_body=response_body_str,
                exception_info=None,
                user_ip=client_ip,
                duration=duration
            )

            background_tasks = BackgroundTasks()
            background_tasks.add_task(self.save_log, log_entry)

            return Response(content=response_content, status_code=response.status_code, headers=dict(response.headers),
                            background=background_tasks)
        except Exception as e:
            print(traceback.format_exc())
            error_body = json.dumps({"error": "Error occurred while processing request"}, ensure_ascii=False)
            end_time = datetime.now(beijing_tz)
            duration = (end_time - start_time_without_ms).total_seconds()
            log_entry = LogEntry(
                timestamp=start_time,
                level="ERROR",
                module=request.url.path,
                function=request.method,
                message=str(e),
                params={**params, **query_params},
                response_status=500,
                response_body=error_body,
                exception_info=traceback.format_exc(),
                user_ip=client_ip,
                duration=duration

            )
            background_tasks = BackgroundTasks()
            background_tasks.add_task(self.save_log, log_entry)
            return Response(content=error_body, status_code=500, headers={'Content-Type': 'application/json'},
                            background=background_tasks)


    @staticmethod
    async def save_log(log_entry):
        await log_entry.save_async()
